# -*- coding: utf-8 -*-
"""
一键把带顶点颜色的3D OBJ按照 Vertex Color → PNG贴图并导出 OBJ+MTL 或 FBX
用法：
    blender your_file.blend  --background --python convert_3d_obj_to_use_texture.py  -- -output_dir /tmp/baked/ -input_obj xxx.obj -out_format obj
"""
import os,sys 
import bpy

# 获取当前文件的绝对路径，并返回上级目录 
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__),  '..'))
sys.path.insert(0,  parent_dir)  # 将上级目录添加到搜索路径首部 
sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))

from common_func import export_model, parse_cli_args, delete_default_objs


# ---------- STEP0:准备 ----------
scene = bpy.context.scene  
scene.render.engine  = 'CYCLES'
scene.cycles.samples  = 32         # bake速度快一点 
 

def bake_vertex_color_to_texture(obj, out_dir):
    # ---------- STEP1:创建/复用材质 ----------
    mat_name = f"{obj.name}_Mat" 
    if mat_name in bpy.data.materials: 
        mat = bpy.data.materials[mat_name] 
    else:
        mat = bpy.data.materials.new(name=mat_name) 
        mat.use_nodes  = True 
    
    obj.data.materials.clear() 
    obj.data.materials.append(mat) 
    
    nodes = mat.node_tree.nodes  
    links = mat.node_tree.links  
    
    # ---------- STEP2:建立节点 ----------
    nodes.clear() 
    out_node = nodes.new(type='ShaderNodeOutputMaterial') 
    em_node = nodes.new(type='ShaderNodeEmission') 
    vc_node = nodes.new(type='ShaderNodeVertexColor') 
    
    # link: VertexColor -> Emission -> Output 
    links.new(vc_node.outputs['Color'],  em_node.inputs['Color']) 
    links.new(em_node.outputs['Emission'],  out_node.inputs['Surface']) 
    
    # ---------- STEP3:UV ----------
    # for o in scene.objects: 
    #     o.select_set(False) 
    
    mesh_objs=bpy.context.selected_objects 
    if not mesh_objs:
        raise RuntimeError('场景中没有 MESH')

    obj=mesh_objs[0]
    view=bpy.context.view_layer  
    view.objects.active=obj    #设为 active 
    obj.select_set(True) 
    
    if not obj.data.uv_layers: 
        obj.data.uv_layers.new(name="UVMap") 
    
    bpy.ops.object.mode_set(mode='EDIT') 
    bpy.context.view_layer.update() 
    try:
        bpy.ops.mesh.select_all(action='SELECT') 
        bpy.ops.uv.smart_project(angle_limit=30*3.14159/180,  island_margin=0.03)
    finally:
        bpy.ops.object.mode_set(mode='OBJECT') 
    
    # ---------- STEP4:创建空白烘焙贴图 ----------
    img_name = f"{obj.name}_VC_baked" 
    if img_name in bpy.data.images: 
        img = bpy.data.images[img_name] 
    else:
        img = bpy.data.images.new(name=img_name, 
                                width=1024,
                                height=1024,
                                alpha=False,
                                float_buffer=False)
    
    # ---------- STEP5:Image Texture节点 ----------
    img_node = nodes.new(type='ShaderNodeTexImage') 
    img_node.image  = img 

   
    # ---------- STEP6:Bake ----------
    scene.cycles.bake_type  = 'EMIT'
    scene.render.bake.margin  = 2           #像素外延防渗色 
    
    # Cycles要求目标ImageTexture节点被激活(蓝色高亮)
    nodes.active  = img_node               #关键一步！
    
    bpy.ops.object.bake(type='EMIT',  use_clear=True)
    
    # ---------- STEP7:保存PNG ----------
    png_path = os.path.join(out_dir,  f"{img_name}.png")
    img.save_render(filepath=png_path) 
    
    img.filepath_raw  = png_path        #指向磁盘绝对路径，这里特别重要，否则无法在保存模型的时候关联到纹理
    img.pack()                          #嵌入 blend(可选)，或在后面 Mark as Local也行 

     # ---------- STEP8:调整导出关联纹理 ----------
    bsdf_node = nodes.new(type='ShaderNodeBsdfPrincipled') 
    links.new(img_node.outputs['Color'],  bsdf_node.inputs['Base Color']) 
    links.new(bsdf_node.outputs['BSDF'],  out_node.inputs['Surface']) 
        
    nodes.remove(vc_node)
    nodes.remove(em_node)
    


if __name__ == "__main__":

# ---------- CLI参数 ----------
    args = parse_cli_args()
    out_dir = args.get("output_dir", os.getcwd())
    input_obj = args.get("input_obj") 
    out_format =  args.get("out_format", "obj") 
    if not input_obj:
        print(f"您需要输入一个3D OBJ模型文件用于转换，输入参数： {args}")
        sys.exit(0)

    if out_format not in ['OBJ', 'obj', 'Obj', 'FBX', 'fbx', 'Fbx']:
        print(f"您输入的导出格式不支持，目前只支持导出OBJ或FBX，请输入其中之一['OBJ', 'obj', 'Obj', 'FBX', 'fbx', 'Fbx']，输入参数： {args}")
        sys.exit(0)


    print(f"您需要将{input_obj}转换成{out_format}，并输出到目录{out_dir}")

    os.makedirs(out_dir,  exist_ok=True)

    delete_default_objs()


    with bpy.context.temp_override( 
            selected_objects=bpy.context.scene.objects[:],    #先清空旧选区保险起见 
            active_object=None):
        for ob in list(bpy.data.objects): 
            ob.select_set(False) 

    with bpy.context.temp_override(filepath=str(input_obj)): 
        bpy.ops.wm.obj_import(filepath=str(input_obj)) 

        # 调整模型的几何中心
        bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY',  center='BOUNDS')

        obj = bpy.context.view_layer.objects.active  or scene.objects[0] 
        assert obj and obj.type  == 'MESH', '请先选中一个网格物体'

        bake_vertex_color_to_texture(obj, out_dir)
        
        # ---------- 导出OBJ+MTL 或 FBX ----------
        model_path = os.path.join(out_dir,  f"{obj.name}.{out_format}") 
        with bpy.context.temp_override(selected_objects=[obj]): 
            export_model(model_path)
        
        print(f"已输出到 {out_dir}")
